#include <MorphManager.h>
#include <InformationManager.h>

MorphManager::MorphManager(Arbitrator::Arbitrator<BWAPI::Unit*,double>* arbitrator)
{
  this->arbitrator = arbitrator;
  for(std::set<BWAPI::UnitType>::iterator i=BWAPI::UnitTypes::allUnitTypes().begin();i!=BWAPI::UnitTypes::allUnitTypes().end();i++)
  {
    plannedCount[*i]=0;
    startedCount[*i]=0;
  }
}

bool MorphManager::canMake(BWAPI::Unit* builder, BWAPI::UnitType type)
{
  if (builder != NULL)
  {
    /* Check if the owner of the unit is you */
    if (builder->getPlayer() != BWAPI::Broodwar->self())
      return false;

    /* Check if this unit can actually build the unit type */
    if (builder->getType() != (type.whatBuilds().first))
      return false;
  }

  for(std::map<BWAPI::UnitType, int>::const_iterator i = type.requiredUnits().begin(); i != type.requiredUnits().end(); i++)
  {
    bool pass = false;
    if (BWAPI::Broodwar->self()->completedUnitCount(i->first) >= i->second)
      pass = true;
    if (i->first == BWAPI::UnitTypes::Zerg_Hatchery)
    {
      if (BWAPI::Broodwar->self()->allUnitCount(BWAPI::UnitTypes::Zerg_Lair) >= i->second)
        pass = true;
      if (BWAPI::Broodwar->self()->allUnitCount(BWAPI::UnitTypes::Zerg_Hive) >= i->second)
        pass = true;
    }
    if (i->first == BWAPI::UnitTypes::Zerg_Lair)
      if (BWAPI::Broodwar->self()->allUnitCount(BWAPI::UnitTypes::Zerg_Hive) >= i->second)
        pass = true;
    if (i->first == BWAPI::UnitTypes::Zerg_Spire)
      if (BWAPI::Broodwar->self()->allUnitCount(BWAPI::UnitTypes::Zerg_Greater_Spire) >= i->second)
        pass = true;
    if (pass == false)
      return false;
  }

  if (type.requiredTech() != BWAPI::TechTypes::None)
    if (!BWAPI::Broodwar->self()->hasResearched((type.requiredTech())))
      return false;

  return true;
}

void MorphManager::onOffer(std::set<BWAPI::Unit*> units)
{
  //go through all the units that are being offered to us
  for(std::set<BWAPI::Unit*>::iterator i=units.begin();i!=units.end();i++)
  {
    //if the unit is a hatchery, make sure either we have no main hatch 
    //  or the offered IS the main hatch
    bool wrongHatch = false;
    if ((*i)->getType() == BWAPI::UnitTypes::Zerg_Hatchery){
      BWAPI::TilePosition homeBase = BWAPI::Broodwar->self()->getStartLocation();
      std::set<BWAPI::Unit*> homeUnits = BWAPI::Broodwar->getUnitsOnTile(homeBase.x(), homeBase.y());
      for (std::set<BWAPI::Unit*>::iterator h=homeUnits.begin(); h!=homeUnits.end(); ++h){
        if ((*h)->getType() == BWAPI::UnitTypes::Zerg_Hatchery){
          // found a hatchery, check if its the same
          if (*h != *i){
            wrongHatch = true;
          }
          break;
        }
      }
    }

    //we will go through all the units in the morph queue
    // if you are supply blocked for at least one unit in the queue, you should rush to Overlord even if you can build a different
    bool supplyBlocked = false; 
    std::map<BWAPI::UnitType,std::list<BWAPI::UnitType> >::iterator q=morphQueues.find((*i)->getType());
    bool used=false;
    // sometimes we are morphing Mutalisks or Hydralisks; so make sure they
    // aren't attacking before we use them
    if (!(*i)->isAttacking()) {
      if (!wrongHatch && q!=morphQueues.end() && !q->second.empty()) //well, only if the queue actually exists
      {
        for(std::list<BWAPI::UnitType>::iterator t=q->second.begin();t!=q->second.end();t++) //loop through the queue
        {
          if (canMake(*i,*t)) //only accept if the given unit can make the type of unit we want to make
          {
            // If we are supply blocked, only use Overlords, otherwise skip
            if (!supplyBlocked || *t == BWAPI::UnitTypes::Zerg_Overlord) {
              if (this->needOverlord(*t)) {
                supplyBlocked = true;
              } else {
                Unit newUnit;
                newUnit.type=*t;
                newUnit.started=false;
                morphingUnits.insert(std::make_pair(*i,newUnit));
                q->second.erase(t);
                //tell the arbitrator we accept the unit, and raise the bid to hopefully prevent other managers from using it
                arbitrator->accept(this,*i);
                arbitrator->setBid(this,*i,100.0);
                used=true;
                break;
              }
            }
          }
          else
          {
            if (BWAPI::Broodwar->self()->minerals()<(*t).mineralPrice() || BWAPI::Broodwar->self()->gas()<(*t).gasPrice())
            {
              break;
            }
          }
        }
      }
    }
    //if we didnt use this unit, tell the arbitrator we decline it
    if (!used)
    {
      arbitrator->decline(this,*i,0);
      arbitrator->removeBid(this,*i);
    }
  }
}

void MorphManager::onRevoke(BWAPI::Unit* unit, double bid)
{
  std::map<BWAPI::Unit*,Unit>::iterator m=morphingUnits.find(unit);
  if (m!=morphingUnits.end())
  {
    BWAPI::UnitType type=m->second.type;
    if (m->first->getType()!=type || m->first->isMorphing())
      morphQueues[type.whatBuilds().first].push_front(type);
    morphingUnits.erase(m);
  }
}

void MorphManager::update()
{
  std::set<BWAPI::Unit*> myPlayerUnits=BWAPI::Broodwar->self()->getUnits();
  for(std::set<BWAPI::Unit*>::iterator u = myPlayerUnits.begin(); u != myPlayerUnits.end(); u++)
  {
    std::map<BWAPI::UnitType,std::list<BWAPI::UnitType> >::iterator q=morphQueues.find((*u)->getType());
    if (q!=morphQueues.end() && !q->second.empty() && (*u)->isCompleted() && morphingUnits.find(*u)==morphingUnits.end())
      arbitrator->setBid(this, *u, 50);
  }
  std::map<BWAPI::Unit*,Unit>::iterator i_next;
  //iterate through all the morphing units
  for(std::map<BWAPI::Unit*,Unit>::iterator i=morphingUnits.begin();i!=morphingUnits.end();i=i_next)
  {
    i_next=i;
    i_next++;
    if (i->first->isCompleted())
    {
      //if the unit is completed and is the right type, we are done
      if (i->first->getType()==i->second.type)
      {
        arbitrator->removeBid(this, i->first);
        plannedCount[i->second.type]--;
        startedCount[i->second.type]--;
        morphingUnits.erase(i);
      }
      else {//if the unit is not the right type, tell it to morph into the right type
        if (BWAPI::Broodwar->canMake(NULL,i->second.type)){
          i->first->morph(i->second.type);
        } else {
          if (this->needOverlord(i->second.type)){
            // revoke and re-offer this larva; onOffer will prefer the overlord to the current unit
            std::set<BWAPI::Unit*> temp;
            temp.insert(i->first);
            this->onRevoke(i->first, 0);
            this->onOffer(temp);
          }
        }
      }
    }
    // check the time: if we start it, it is considered morphing even if it will fail?
    else if (i->first->isMorphing() && i->first->getRemainingBuildTime() < i->second.type.buildTime())
    {
      //if the unit is morphing, note that we have started it in the startedCount
      if (!i->second.started)
      {
        startedCount[i->second.type]++;
        i->second.started=true;

        // record spending
        TheInformationManager->removeItemOrder(i->second.type.getName());
      }
    }
  }
}

std::string MorphManager::getName() const
{
  return "Morph Manager";
}

std::string MorphManager::getShortName() const
{
  return "Morph";
}

void MorphManager::onRemoveUnit(BWAPI::Unit* unit)
{
  std::map<BWAPI::Unit*,Unit>::iterator m=morphingUnits.find(unit);
  if (m!=morphingUnits.end())
  {
    BWAPI::UnitType type=m->second.type;
    morphQueues[type.whatBuilds().first].push_front(type);
    morphingUnits.erase(m);
  }
}

bool MorphManager::morph(BWAPI::UnitType type)
{
  //morph order starts here

  //we only accept buildings that are made from other buildings
  //and units that are made from other units

  if (type.isBuilding()!=type.whatBuilds().first.isBuilding()) 
    return false;
  morphQueues[type.whatBuilds().first].push_back(type);
  plannedCount[type]++;
  return true;
}

int MorphManager::getPlannedCount(BWAPI::UnitType type) const
{
  std::map<BWAPI::UnitType,int>::const_iterator i=plannedCount.find(type);
  if (i!=plannedCount.end())
    return i->second;
  return 0;
}

int MorphManager::getStartedCount(BWAPI::UnitType type) const
{
  std::map<BWAPI::UnitType,int>::const_iterator i=startedCount.find(type);
  if (i!=startedCount.end())
    return i->second;
  return 0;
}
BWAPI::UnitType MorphManager::getBuildType(BWAPI::Unit* unit) const
{
  if (morphingUnits.find(unit)==morphingUnits.end())
    return BWAPI::UnitTypes::None;
  return morphingUnits.find(unit)->second.type;
}

bool MorphManager::needOverlord(BWAPI::UnitType t)
{
  // check supply
  if (t.supplyRequired() > 0 && t.supplyRequired() > BWAPI::Broodwar->self()->supplyTotal() - BWAPI::Broodwar->self()->supplyUsed() && BWAPI::Broodwar->self()->supplyTotal() < 400){
    // check for overlord in queue
    if (this->getPlannedCount(BWAPI::UnitTypes::Zerg_Overlord) > 0) {
      return true;
    }
  }
  return false;
}
